package com.kara.woodAgent.agent.graph;

import com.kara.woodAgent.agent.graph.communicate.*;
import com.kara.woodAgent.agent.graph.exception.InnerCheckException;
import com.kara.woodAgent.agent.graph.impl.ResumePayload;
import com.kara.woodAgent.agent.wood.playload.ThinksPayload;
import com.kara.woodAgent.common.exception.BizException;

import java.util.Objects;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author tzy
 * @version 1.0
 * @project wood_agent
 * @description 代表当前执行到的步骤
 * @date 2025/6/4 10:50:28
 */
public class BaseRunningGraph<T extends GraphContext> implements RunningGraph<T> {


    private String nodeName = "__init__";

    private final BaseGraph<T> baseGraph;

    public Consumer<String> consumer;

    public BiConsumer<String, ThinksPayload> function;

    private final T context;

    public final Channel channel;

    // init  0
    // running 1
    // stop 2
    // checking 3

    public static int STATUS_INIT = 0;
    public static int STATUS_RUNNING = 1;
    public static int STATUS_FINISHED = 2;
    public static int STATUS_CHECKING = 3;

    private int status = 0;

    private StepImage<T> stepImage;

    private Thread runningThread;

    public BaseRunningGraph(String nodeName, BaseGraph<T> baseGraph, T context, Channel channel) {
        this.baseGraph = baseGraph;
        this.context = context;
        this.channel = channel;
        this.nodeName = nodeName;
        //注册checkPoint事件
        this.channel.addCheckpointListener(payload -> {
            //如果已经生成了镜像,纪实从检查点恢复模式，直接返回镜像
            if (payload instanceof ResumePayload) {

                if (this.status == STATUS_CHECKING && stepImage != null) {

                    ChannelMessage channelMessage = getChannelMessage();
                    this.stepImage = null;
                    this.status = STATUS_RUNNING;
                    return channelMessage;
                } else {
                    throw new RuntimeException("graph status not checking, can not resume");
                }
            }
            this.status = STATUS_CHECKING;
            //生成镜像
            makeImage(payload);
            //用于内部终端执行
            throw new InnerCheckException();
        });

        this.channel.addStatusListener(() -> this.status == STATUS_CHECKING);
    }

    private ChannelMessage getChannelMessage() {
        ChannelMessage channelMessage = this.stepImage.getChannelMessage();

        //子图触发 需要返回子图的镜像
        if (stepImage.subGraphCheck() && stepImage.getChannelPayload() instanceof TransmitPayload transmitPayload) {
            StepImage<T> subImage = (StepImage<T>) transmitPayload.getImage();
            subImage.setChannelMessage(channelMessage);
            channelMessage = new TransmitMessage<>(subImage);
//            channelMessage.setContent(content);
        }
        return channelMessage;
    }

    private void makeImage(ChannelPayload payload) {


        stepImage = new StepImage<>();

        stepImage.setNode(nodeName);
        stepImage.setContext(context);
        stepImage.setGraphClassName(baseGraph.getClass().getName());
        stepImage.setChannelPayload(payload);
        stepImage.setStatus(payload instanceof TransmitPayload ? StepImage.STATUS_SUB_GRAPH_CHECKING : StepImage.STATUS_CHECKING);

    }

    public void saveImage(StepImage<T> image) {
        this.stepImage = image;
        this.status = Objects.equals(image.getStatus(), StepImage.STATUS_CHECKING) || Objects.equals(image.getStatus(), StepImage.STATUS_SUB_GRAPH_CHECKING) ? STATUS_CHECKING : STATUS_INIT;
    }


    public void nodeName(String nodeName) {
        this.nodeName = nodeName;
    }

    public String nodeName() {
        return nodeName;
    }


    public void execute(String response) {
        if (this.consumer != null) {
            this.consumer.accept(response);
        }
    }

    public void start() {
        if(status!=STATUS_INIT && status!=STATUS_CHECKING){
            throw new BizException("graph status not init or checking, can not start");
        }
        if(this.status==STATUS_INIT){
            this.status = STATUS_RUNNING;
        }
        this.runningThread = Thread.currentThread();
        this.baseGraph.execute(context, this);
    }

    @Override
    public void stop() {
        this.status = STATUS_FINISHED;
        if (this.runningThread != null) {
            this.runningThread.interrupt();
        }
    }

    @Override
    public boolean isStop() {
        if (this.runningThread == null) {
            return true;
        }
        return this.runningThread.isInterrupted();
    }

    @Override
    public void clear() {
        if (this.status == STATUS_FINISHED) {
            this.runningThread = null;
        }
    }

    @Override
    public void finish() {
        this.status = STATUS_FINISHED;
    }

    @Override
    public boolean checking() {
        return this.status == STATUS_CHECKING;
    }

    @Override
    public int status() {
        return status;
    }

    @Override
    public StepImage<T> image() {
        if (!checking() || stepImage == null) {
            throw new RuntimeException("don't have image, please checkpoint first");
        }
        return stepImage;
    }

    @Override
    public RunningGraph<T> then(Consumer<String> consumer) {
        this.consumer = consumer;
        return this;
    }

    @Override
    public RunningGraph<T> communicate(BiConsumer<String, ThinksPayload> function) {
        this.function = function;
        channel.addWaitUserInputEventListener(payload -> this.function.accept(nodeName, payload));
        return this;
    }

    @Override
    public RunningGraph<T> listen(String event, Function<ChannelPayload, ChannelMessage> func) {
        channel.addEventListener(event, func);
        return this;
    }
}
